# 标准库导入
import os  # 操作系统相关功能
import argparse  # 命令行参数解析
import pandas as pd  # 数据处理和分析

# 类型提示用于更好的代码文档
from typing import Dict, List, Tuple, Optional

# 本地模块导入 - 传感器分类系统
from data_loader import prepare_dataset  # 加载和准备传感器数据
from feature_extraction import extract_features_from_segments, normalize_features  # 特征工程
from model import SensorClassifier  # 机器学习模型类
from train import train_model  # 模型训练功能
from predict import predict_directory, predict_single_file  # 预测功能
from visualize import visualize_data  # 数据可视化功能
from real_time_predict import predict_real_time  # 实时预测模拟
from data_preparation import prepare_data_files, get_available_classes  # 数据准备工具
from mqtt_predict import predict_with_mqtt  # 基于MQTT的实时预测
from influxdb_loader import data_with_influxdb  #从influxdb获取数据
from execute_shellscripts import exec_scripts #从mysql读取脚本并执行


def main():
    """物联网传感器分类系统的主入口点。
    处理命令行参数并执行请求的命令。"""
    
    # 获取项目根目录的绝对路径
    script_dir = os.path.dirname(os.path.abspath(__file__))  # 包含此脚本的目录
    project_root = os.path.dirname(script_dir)  # 父目录是项目根目录
    
    # 配置相对于项目根的默认路径
    default_data_dir = os.path.join(project_root, "data")  # 默认数据目录
    default_output_dir = os.path.join(project_root, "output")  # 默认输出目录
    default_model_path = os.path.join(default_output_dir, "model.joblib")  # 默认模型保存路径
    default_scaler_path = os.path.join(default_output_dir, "scaler.joblib")  # 默认特征缩放器保存路径
    default_vis_dir = os.path.join(project_root, "visualizations")  # 默认可视化输出目录
    
    # 设置带有子命令的主参数解析器
    parser = argparse.ArgumentParser(description='物联网传感器分类系统')
    subparsers = parser.add_subparsers(dest='command', required=True,
                                     help='运行系统的可用命令')
    
    # 获取数据 - 从influxdb获取训练数据
    influxdb_parser = subparsers.add_parser('influxdb',help='获取训练数据')
    influxdb_parser.add_argument('--bucket', type=str, required=True, help='库名称')
    influxdb_parser.add_argument('--measurement',type=str,required=True,help='表名称')
    influxdb_parser.add_argument('--deviceid',type=str,required=True,help='设备Mac地址')
    influxdb_parser.add_argument('--stime',type=str,required=True,help='数据开始时间')
    influxdb_parser.add_argument('--etime',type=str,required=True,help='数据结束时间')

    # 获取执行脚本 - 从Mysql获取执行脚本
    exec_parser = subparsers.add_parser('execshell',help='获取执行脚本数据')
    exec_parser.add_argument('--status', type=str,default='00000',help='数据状态')
    
    # 训练命令 - 训练机器学习模型
    train_parser = subparsers.add_parser('train', help='训练分类模型')
    train_parser.add_argument('--data_dir', type=str, default=default_data_dir,
                             help='包含训练数据文件的目录路径')
    train_parser.add_argument('--model_type', type=str, default='rf',
                             choices=['rf', 'svm'], help='模型类型: rf=随机森林, svm=支持向量机')
    train_parser.add_argument('--window_size', type=int, default=100,
                             help='滑动窗口大小(样本数, 默认: 100)')
    train_parser.add_argument('--step_size', type=int, default=50,
                             help='窗口之间的步长(样本数, 默认: 50)')
    train_parser.add_argument('--output_dir', type=str, default=default_output_dir,
                             help='保存训练模型和缩放器的目录(默认: ./output)')
    
    # 预测命令
    predict_parser = subparsers.add_parser('predict', help='进行预测')
    predict_parser.add_argument('--data', type=str, required=True,
                               help='要预测的文件或目录路径')
    predict_parser.add_argument('--model', type=str, default=default_model_path,
                               help='训练好的模型路径')
    predict_parser.add_argument('--scaler', type=str, default=default_scaler_path,
                               help='特征缩放器路径')
    predict_parser.add_argument('--window_size', type=int, default=100,
                               help='窗口大小(样本数)')
    predict_parser.add_argument('--step_size', type=int, default=50,
                               help='滑动窗口步长')
    predict_parser.add_argument('--output', type=str, default=None,
                               help='保存预测结果的路径')
    
    # 可视化命令
    visualize_parser = subparsers.add_parser('visualize', help='可视化数据')
    visualize_parser.add_argument('--data_dir', type=str, default=default_data_dir,
                                 help='包含数据文件的目录路径')
    visualize_parser.add_argument('--output_dir', type=str, default=default_vis_dir,
                                 help='保存图表的目录')
    
    # 实时预测命令
    realtime_parser = subparsers.add_parser('realtime', help='实时预测模拟')
    realtime_parser.add_argument('--model', type=str, default=default_model_path,
                               help='训练好的模型路径')
    realtime_parser.add_argument('--scaler', type=str, default=default_scaler_path,
                               help='特征缩放器路径')
    realtime_parser.add_argument('--data_dir', type=str, default=default_data_dir,
                               help='包含样本数据文件的目录')
    realtime_parser.add_argument('--window_size', type=int, default=100,
                               help='窗口大小(样本数)')
    realtime_parser.add_argument('--interval', type=float, default=1.0,
                               help='预测之间的时间间隔(秒)')
    realtime_parser.add_argument('--num_predictions', type=int, default=10,
                               help='要进行的预测次数(默认: 10, 0表示无限)')
    
    # MQTT预测命令
    mqtt_parser = subparsers.add_parser('mqtt', help='基于MQTT的实时预测')
    mqtt_parser.add_argument('--model', type=str, default=default_model_path,
                           help='训练好的模型路径')
    mqtt_parser.add_argument('--scaler', type=str, default=default_scaler_path,
                           help='特征缩放器路径')
    mqtt_parser.add_argument('--broker', type=str, required=True,
                           help='MQTT代理地址')
    mqtt_parser.add_argument('--port', type=int, default=1883,
                           help='MQTT代理端口')
    mqtt_parser.add_argument('--topic', type=str, default='sensor/data',
                           help='要订阅的MQTT主题')
    mqtt_parser.add_argument('--username', type=str, default=None,
                           help='MQTT认证用户名')
    mqtt_parser.add_argument('--password', type=str, default=None,
                           help='MQTT认证密码')
    mqtt_parser.add_argument('--window-size', type=int, default=100,
                           help='窗口大小(样本数)')
    mqtt_parser.add_argument('--interval', type=float, default=1.0,
                           help='预测之间的时间间隔(秒)')
    
    # 准备数据命令
    prepare_parser = subparsers.add_parser('prepare', help='准备训练数据文件')
    prepare_parser.add_argument('--data_dir', type=str, default=default_data_dir,
                                help='包含原始数据文件的目录路径')
    prepare_parser.add_argument('--output_dir', type=str, default=os.path.join(project_root, "data_prepared"),
                                help='保存准备数据文件的目录路径')
    
    args = parser.parse_args()
    
    if args.command == 'train':
        train_model(
            data_dir=args.data_dir,
            model_type=args.model_type,
            window_size=args.window_size,
            step_size=args.step_size,
            output_dir=args.output_dir
        )
    
    elif args.command == 'predict':
        # 检查数据路径是文件还是目录
        if os.path.isfile(args.data):
            result = predict_single_file(
                args.data, args.model, args.scaler, args.window_size, args.step_size
            )
            print(f"预测类别: {result['predicted_class']}")
            print("类别概率:")
            for cls, prob in result['average_probabilities'].items():
                print(f"  {cls}: {prob:.4f}")
        
        elif os.path.isdir(args.data):
            results_df = predict_directory(
                args.data, args.model, args.scaler, args.window_size, args.step_size, args.output
            )
            print("\n预测摘要:")
            print(results_df[['filename', 'predicted_class']].to_string(index=False))
        
        else:
            print(f"错误: {args.data} 不是有效的文件或目录")
    
    elif args.command == 'visualize':
        visualize_data(args.data_dir, args.output_dir)
    
    elif args.command == 'realtime':
        # 将0转换为None表示无限预测
        num_predictions = None if args.num_predictions == 0 else args.num_predictions
        
        predict_real_time(
            model_path=args.model,
            scaler_path=args.scaler,
            data_dir=args.data_dir,
            window_size=args.window_size,
            interval=args.interval,
            num_predictions=num_predictions
        )
    
    elif args.command == 'mqtt':
        predict_with_mqtt(
            model_path=args.model,
            scaler_path=args.scaler,
            broker_address=args.broker,
            broker_port=args.port,
            topic=args.topic,
            username=args.username,
            password=args.password,
            window_size=args.window_size,
            interval=args.interval
        )
    
    elif args.command == 'prepare':
        # 获取可用类别
        classes = get_available_classes(args.data_dir)
        print(f"找到 {len(classes)} 个类别: {classes}")
        
        # 准备数据文件
        prepare_data_files(args.data_dir, args.output_dir)

    elif args.command == 'influxdb':
        #从influxdb获取数据
        data_with_influxdb(args.bucket,args.measurement,args.deviceid,args.stime,args.etime)
    
    elif args.command == 'execshell':
        #从mysql获取脚本并执行更新
        exec_scripts()

    else:
        parser.print_help()


if __name__ == "__main__":
    main()
